{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "75e10d3b-9b26-4bca-adf7-3f16e5315bec",
   "metadata": {},
   "source": [
    "Milvus Demo\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2f0fa062-10c6-49ff-8bf5-45b9a2f9dced",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: pymilvus\n",
      "Version: 2.3.0\n",
      "Summary: Python Sdk for Milvus\n",
      "Home-page: \n",
      "Author: \n",
      "Author-email: Milvus Team <milvus-team@zilliz.com>\n",
      "License: \n",
      "Location: /opt/conda/lib/python3.11/site-packages\n",
      "Requires: environs, grpcio, pandas, protobuf, ujson\n",
      "Required-by: \n"
     ]
    }
   ],
   "source": [
    "# Show pymilvus information\n",
    "!pip show pymilvus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "588f2ae0-7f09-436c-8964-01f2ce8ed0c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set Milvus host and port\n",
    "MILVUS_HOST = \"milvus\"  # this is the service name defined in compose yaml\n",
    "MILVUS_PORT = \"19530\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eeaddf68-1441-4ef3-bbc7-beb54a45f937",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import dependencies\n",
    "\n",
    "from pymilvus import connections,Collection, CollectionSchema, FieldSchema, DataType, utility\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ebdebac2-fe07-4fd3-a799-2074a3eae296",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to milvus\n",
    "\n",
    "connections.connect(\"default\", host=MILVUS_HOST, port=MILVUS_PORT)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cbf7cff6-11ac-4e30-8a76-3bdb663a78b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create schema\n",
    "# Please refer to https://milvus.io/docs/create_collection.md#Prepare-Schema\n",
    "\n",
    "book_id = FieldSchema(\n",
    "  name=\"book_id\",\n",
    "  dtype=DataType.INT64,\n",
    "  is_primary=True,\n",
    ")\n",
    "book_name = FieldSchema(\n",
    "  name=\"book_name\",\n",
    "  dtype=DataType.VARCHAR,\n",
    "  max_length=200,\n",
    "  # The default value will be used if this field is left empty during data inserts or upserts.\n",
    "  # The data type of `default_value` must be the same as that specified in `dtype`.\n",
    "  default_value=\"Unknown\"\n",
    ")\n",
    "word_count = FieldSchema(\n",
    "  name=\"word_count\",\n",
    "  dtype=DataType.INT64,\n",
    "  # The default value will be used if this field is left empty during data inserts or upserts.\n",
    "  # The data type of `default_value` must be the same as that specified in `dtype`.\n",
    "  default_value=9999\n",
    ")\n",
    "book_intro = FieldSchema(\n",
    "  name=\"book_intro\",\n",
    "  dtype=DataType.FLOAT_VECTOR,\n",
    "  dim=2\n",
    ")\n",
    "schema = CollectionSchema(\n",
    "  fields=[book_id, book_name, word_count, book_intro],\n",
    "  description=\"Test book search\",\n",
    "  enable_dynamic_field=True\n",
    ")\n",
    "collection_name = \"book\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "35e88b99-40b1-4d2d-8d00-81efb9f7b28f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a collection with the schema\n",
    "# Please refer to https://milvus.io/docs/create_collection.md#Create-a-collection-with-the-schema\n",
    "\n",
    "collection = Collection(\n",
    "    name=collection_name,\n",
    "    schema=schema,\n",
    "    using='default',\n",
    "    shards_num=2\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dd60684c-61b7-4968-baf4-11aba43a6d60",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [[0.9899460808567317, 0.9610429989947455], [0.7866622633345994, 0.23524804243295305], [0.4399626677948054, 0.4236408323595152], [0.6576079177287327, 0.8954567429888496], [0.6224898353815375, 0.678062362078302], [0.3878256222454457, 0.5014967593371562], [0.49983178616582136, 0.3329611266131416], [0.021130813764533296, 0.6545340534345481], [0.8058905227229882, 0.7018186158265323], [0.16794665577145507, 0.15648383157427537]]]\n"
     ]
    }
   ],
   "source": [
    "# Prepare data\n",
    "# Please refer to https://milvus.io/docs/insert_data.md#Prepare-data\n",
    "\n",
    "import random\n",
    "data = [\n",
    "  [i for i in range(10)],\n",
    "  [str(i) for i in range(10)],\n",
    "  [i for i in range(10, 20)],\n",
    "  [[random.random() for _ in range(2)] for _ in range(10)],\n",
    "]\n",
    "\n",
    "print(data, flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a47ecf18-3360-48ca-82af-ff148dcd259d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Insert data\n",
    "# Please refer to https://milvus.io/docs/insert_data.md#Insert-data-to-Milvus\n",
    "\n",
    "collection = Collection(\"book\")      # Get an existing collection.\n",
    "mr = collection.insert(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6d03ad22-09b2-430c-b0af-34c7eaf751fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Status(code=0, message=)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create an index to avoid issue in loading collection\n",
    "\n",
    "index_params = {\n",
    "    \"index_type\": \"IVF_FLAT\",\n",
    "    \"metric_type\": \"L2\",\n",
    "    \"params\": {\n",
    "        \"nlist\": 1\n",
    "    }\n",
    "}\n",
    "\n",
    "collection.create_index(\n",
    "  field_name=\"book_intro\",\n",
    "  index_params=index_params,\n",
    "  index_name=\"book_intro_index\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "97bd2eda-09c0-483d-b519-e1b70cce4c3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data before query from collection\n",
    "# https://milvus.io/docs/query.md#Conduct-a-Query\n",
    "\n",
    "collection = Collection(\"book\")      # Get an existing collection.\n",
    "collection.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b5d8236c-5cc4-483f-bd81-a925c6eeda90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'book_id': 2, 'book_intro': [0.43996266, 0.42364085]}, {'book_id': 4, 'book_intro': [0.6224898, 0.6780624]}, {'book_id': 6, 'book_intro': [0.4998318, 0.3329611]}, {'book_id': 8, 'book_intro': [0.8058905, 0.70181865]}]\n"
     ]
    }
   ],
   "source": [
    "# Run the query\n",
    "\n",
    "res = collection.query(\n",
    "  expr = \"book_id in [2,4,6,8]\",\n",
    "  offset = 0,\n",
    "  limit = 10, \n",
    "  output_fields = [\"book_id\", \"book_intro\"],\n",
    ")\n",
    "\n",
    "sorted_res = sorted(res, key=lambda k: k['book_id'])\n",
    "print(sorted_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0981a889-862e-4163-a946-8a0ab28a82d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up\n",
    "# https://milvus.io/docs/drop_collection.md#Drop-a-collection\n",
    "\n",
    "utility.drop_collection(\"book\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6c57199d-78b7-40e5-b998-fe9ee72ce830",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disconnect from milvus\n",
    "\n",
    "connections.disconnect(\"default\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
